

#include "DriverMain.h"


PDEVICE_OBJECT pDeviceObject = NULL;

VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject);

NTSTATUS DispatchCreate(IN PDEVICE_OBJECT pDeviceObj,
						IN PIRP pIrp);

NTSTATUS DispatchClose(IN PDEVICE_OBJECT pDeviceObj,
						IN PIRP pIrp);

NTSTATUS DispatchRead(IN PDEVICE_OBJECT pDeviceObj,
					   IN PIRP pIrp);

NTSTATUS DispatchWrite(IN PDEVICE_OBJECT pDeviceObj,
					   IN PIRP pIrp);

NTSTATUS DispatchIoControl(IN PDEVICE_OBJECT pDeviceObj,
					       IN PIRP pIrp);

#define  DELAY_ONE_MICROSECOND  (-10)
#define  DELAY_ONE_MILLISECOND (DELAY_ONE_MICROSECOND*1000)
#define  DELAY_ONE_SECOND (DELAY_ONE_MILLISECOND*1000)


VOID KeSleep(int nSecond)
{
	LARGE_INTEGER interval;
	interval.QuadPart = (nSecond*1000 * DELAY_ONE_MILLISECOND);        
	KeDelayExecutionThread(KernelMode,FALSE,&interval);
}


NTSTATUS DriverEntry(IN PDRIVER_OBJECT pDriverObject,
					IN PUNICODE_STRING pRegistryPath)
{
	KdPrint(("Driver Entry \n"));
	KdBreakPoint();
	
	UNICODE_STRING deviceName;
	RtlInitUnicodeString(&deviceName, DEVICE_NAME);
	NTSTATUS status = IoCreateDevice(pDriverObject, sizeof(DRIVER_EXTENSION), &deviceName, 
		FILE_DEVICE_UNKNOWN, NULL, NULL, &pDeviceObject);
	
	if (!NT_SUCCESS(status))
	{
		KdPrint(("IoCreateDevice Failure!"));
		return STATUS_UNSUCCESSFUL;
	}

	UNICODE_STRING symDeviceName;
	RtlInitUnicodeString(&symDeviceName, SYMBOL_DEVICE_NAME);

	status = IoCreateSymbolicLink(&symDeviceName, &deviceName); 
	if (!NT_SUCCESS(status))
	{
		IoDeleteDevice(pDeviceObject);
		pDeviceObject = NULL;
		KdPrint(("IoCreateSymbolicLink Failure!"));
		return STATUS_UNSUCCESSFUL;
	}

	pDeviceObject->Flags |= DO_DIRECT_IO;

	pDriverObject->DriverUnload = DriverUnload;
	pDriverObject->MajorFunction[IRP_MJ_CREATE] = DispatchCreate;
	pDriverObject->MajorFunction[IRP_MJ_CLOSE]  = DispatchClose;
	pDriverObject->MajorFunction[IRP_MJ_READ]   = DispatchRead;
	pDriverObject->MajorFunction[IRP_MJ_WRITE]   = DispatchWrite;
	pDriverObject->MajorFunction[IRP_MJ_DEVICE_CONTROL]   = DispatchIoControl;

	return status;
}

VOID DriverUnload(IN PDRIVER_OBJECT pDriverObject)
{
	UNICODE_STRING symbolDeviceName = RTL_CONSTANT_STRING(SYMBOL_DEVICE_NAME);
	IoDeleteSymbolicLink(&symbolDeviceName);
	if (pDeviceObject)
	{
		IoDeleteDevice(pDeviceObject);
		pDeviceObject = NULL;
	}
	KdPrint(("Driver Unload \n"));
}



// CreateFile
NTSTATUS DispatchCreate(IN PDEVICE_OBJECT pDeviceObj,
						IN PIRP pIrp)
{
	KdPrint(("DispatchCreate"));
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}


// CloseHandle
NTSTATUS DispatchClose(IN PDEVICE_OBJECT pDeviceObj,
					   IN PIRP pIrp)
{
	KdPrint(("DispatchClose")); 
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}



// ReadFile
NTSTATUS DispatchRead(IN PDEVICE_OBJECT pDeviceObj,
					  IN PIRP pIrp)
{
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	ULONG len = stack->Parameters.Read.Length;  // ReadFile data length
	LARGE_INTEGER offset = stack->Parameters.Read.ByteOffset;  // ReadFile offset;
	KdPrint(("DispatchRead len = %d", len));
	
	
	// DO_DIRECT_IO
	char* pbuf = (char*) MmGetSystemAddressForMdlSafe(pIrp->MdlAddress, NormalPagePriority);
	if (pbuf)
	{
		// do something
	}

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = len;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}


// WriteFile
NTSTATUS DispatchWrite(IN PDEVICE_OBJECT pDeviceObj,
					   IN PIRP pIrp)
{
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	ULONG len = stack->Parameters.Write.Length;  // WriteFile data length
	LARGE_INTEGER offset = stack->Parameters.Write.ByteOffset;  // WriteFile offset;
	KdPrint(("DispatchWrite len = %d", len));

	// DO_DIRECT_IO
	char* pbuf = (char*) MmGetSystemAddressForMdlSafe(pIrp->MdlAddress, NormalPagePriority);
	if (pbuf)
	{
		// do something
	}
	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = len;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}

NTSTATUS MemTest()
{
	PWCHAR pstr = (PWCHAR) ExAllocatePoolWithTag(NonPagedPool, 1024, MEM_TAG);
	UNICODE_STRING uniStr = {};
	RtlInitEmptyUnicodeString(&uniStr, pstr, 1024);
	RtlUnicodeStringPrintf(&uniStr, L"Hello ProcessId = %d", PsGetCurrentProcessId());
	KdPrint(("%wZ", &uniStr));
	ExFreePoolWithTag(pstr, MEM_TAG);
	return STATUS_SUCCESS;
}

NTSTATUS TestProcess()
{
	ULONG MajVer = 0;
	ULONG MinVer = 0;
	ULONG BuildNum = 0;
	PsGetVersion(&MajVer, &MinVer, &BuildNum, NULL);
	PWCHAR buf = (PWCHAR) ExAllocatePoolWithTag(NonPagedPool, 1024, MEM_TAG);
	RtlZeroMemory(buf, 1024);
	RtlStringCchPrintfW(buf, 1024, L"MajorVersion = %d, MinjorVersion = %d, BuildNumber = %d", MajVer, MinVer, BuildNum);
	
	KdPrint(("%S", buf));
	ExFreePoolWithTag(buf, MEM_TAG);


	HANDLE pidd = PsGetCurrentProcessId();
	KdPrint(("Pid = %d", pidd));

	
	return STATUS_SUCCESS;
}

NTSTATUS FileTest()
{
	HANDLE fileHandle = NULL;
	UNICODE_STRING fileName = RTL_CONSTANT_STRING(L"\\??\\C:\\xxxx\\a.txt");

	OBJECT_ATTRIBUTES objAttr = {};
	InitializeObjectAttributes(&objAttr, &fileName, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE,NULL, NULL);
	IO_STATUS_BLOCK ioStatus = {};
	NTSTATUS statue = ZwCreateFile(&fileHandle, GENERIC_READ | GENERIC_WRITE, &objAttr, &ioStatus, 
		NULL, FILE_ATTRIBUTE_NORMAL, FILE_SHARE_READ, FILE_OPEN_IF, 
		FILE_NON_DIRECTORY_FILE | FILE_RANDOM_ACCESS | FILE_SYNCHRONOUS_IO_NONALERT, NULL, 0);

	if (NT_SUCCESS(statue) && NT_SUCCESS(ioStatus.Status)) 
	{
		KdPrint(("CreateFile Success"));
		char* buf = (char*)ExAllocatePoolWithTag(NonPagedPool, 1024, MEM_TAG);
		RtlStringCchPrintfA(buf, 1024, "Hello Driver pid = %d", PsGetCurrentProcessId());
		ZwWriteFile(fileHandle, NULL, NULL, NULL, &ioStatus, buf, 1024, NULL, NULL);
		ExFreePoolWithTag(buf, MEM_TAG);
		ZwClose(fileHandle);
	}
	else 
	{
		KdPrint(("CreateFile Failure!"));
	}
	return STATUS_SUCCESS;
}


NTSTATUS DirectoryTest()
{
	HANDLE hDir = NULL;
	UNICODE_STRING dirName = RTL_CONSTANT_STRING(L"\\Device");
	OBJECT_ATTRIBUTES objAttr = {};
	InitializeObjectAttributes(&objAttr, &dirName, OBJ_CASE_INSENSITIVE | OBJ_KERNEL_HANDLE, NULL, NULL);

	KdPrint(("%wZ", &dirName));
	NTSTATUS status = ZwCreateDirectoryObject(&hDir, DIRECTORY_ALL_ACCESS, &objAttr);
	if (NT_SUCCESS(status))
	{
		KdPrint(("ZwCreateDirectoryObject Success"));
	}
	else 
	{
		KdPrint(("ZwCreateDirectoryObject Failure"));
	}

	ZwClose(hDir);
	return status;
}

VOID DpcRoutine(   __in struct _KDPC *Dpc,
				   __in_opt PVOID DeferredContext,
				   __in_opt PVOID SystemArgument1,
				   __in_opt PVOID SystemArgument2
				   )
{
	KdPrint(("222222"));
	KdPrint(("Hello this DPC"));
}

VOID DpcTest()
{
	KTIMER timer = {};
	KDPC dpc  = {};
	KeInitializeDpc(&dpc, &DpcRoutine, NULL);
	LARGE_INTEGER dueTime = RtlConvertLongToLargeInteger(-5000000);
	KeInitializeTimer(&timer);
	KeSetTimer(&timer, dueTime, &dpc);
	KdPrint(("111111"));
	KeWaitForSingleObject(&timer, Executive, KernelMode, FALSE, NULL);
	KdPrint(("Wait finished."));
	
}

VOID KThreadProc(IN PVOID context)
{
	for (int i = 0; i < 5; i++)
	{
		KdPrint(("KThreadProc %d", (i)));
		KeSleep(2);
	}
	KdPrint(("Exit Thread"));
	PsTerminateSystemThread(STATUS_SUCCESS); 
}

VOID ThreadTest()
{
	KdPrint(("ThreadTest"));
	HANDLE hThread = NULL;
	OBJECT_ATTRIBUTES oa = {};
	InitializeObjectAttributes(&oa, NULL, OBJ_KERNEL_HANDLE, NULL, NULL);
	CLIENT_ID cid = {};
	NTSTATUS status = PsCreateSystemThread(&hThread, GENERIC_READ | GENERIC_WRITE,
		&oa, NtCurrentProcess(), &cid, KThreadProc, NULL);

	
	KdPrint(("status = %8X, hThread = %d", status , hThread));
	ZwClose(hThread);
}


VOID EventTest()
{
	KEVENT ev;
	KeInitializeEvent(&ev, SynchronizationEvent, FALSE);

	KeSetEvent(&ev, IO_NO_INCREMENT, TRUE);

	KeResetEvent(&ev);

}

// DeviceIoControl
NTSTATUS DispatchIoControl(IN PDEVICE_OBJECT pDeviceObj,
						   IN PIRP pIrp)
{
	PIO_STACK_LOCATION stack = IoGetCurrentIrpStackLocation(pIrp);
	ULONG inLen = stack->Parameters.DeviceIoControl.InputBufferLength;
	ULONG outLen = stack->Parameters.DeviceIoControl.OutputBufferLength;
	ULONG ioCode = stack->Parameters.DeviceIoControl.IoControlCode;
	PVOID inBuffer = NULL;

	KdPrint(("ioCode = %d inLen = %d outLen = %d", ioCode, inLen, outLen));
	

	switch (ioCode)
	{
	case IOCTROL_CODE1:
		// Buffered IO
		inBuffer = pIrp->AssociatedIrp.SystemBuffer;
		break;

	case IOCTL_PROCESS_TEST:
		TestProcess();
		FileTest();
		break;


	case IOCTL_DPC_TEST:
		DpcTest();
		break;

	case IOCTL_THREAD_TEST:
		ThreadTest();
		break;
	}



	//
	//
	//DirectoryTest();

	pIrp->IoStatus.Status = STATUS_SUCCESS;
	pIrp->IoStatus.Information = outLen;
	IoCompleteRequest(pIrp, IO_NO_INCREMENT);
	return STATUS_SUCCESS;
}